#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
A股股票智能筛选系统 v2.0
优化版本：基于实时可靠数据源进行筛选
"""

import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import warnings
import time
warnings.filterwarnings('ignore')

# 配置参数
class ScreenerConfig:
    # 基础条件
    MIN_PRICE = 7
    MAX_PRICE = 50
    MIN_VOLUME = 50_000_000  # 5000万元成交额
    MIN_LIST_YEARS = 2
    MIN_MARKET_CAP = 50_00_000_000  # 50亿元
    
    # 基本面指标（基于实时数据）
    MIN_PE = 5
    MAX_PE = 30
    MIN_PB = 0.8
    MAX_PB = 5
    
    # 技术面指标
    MIN_3M_RETURN = -10  # 百分比
    MAX_3M_RETURN = 20
    MAX_52W_DRAWDOWN = 30  # 百分比
    MIN_TURNOVER = 0.5  # 最小换手率%
    
    # 优先行业
    PRIORITY_INDUSTRIES = [
        '半导体', '集成电路', '芯片', '软件', '云计算', '人工智能', 'AI', '计算机',
        '创新药', 'CXO', '医疗器械', '医疗服务', '生物医药', '医药',
        '锂电', '光伏', '储能', '新能源', '电池', '风电',
        '自动化', '精密', '新材料', '军工',
        '消费', '食品', '饮料', '白酒', '零售'
    ]
    
    OUTPUT_COUNT = 20
    MAX_STOCKS_TO_PROCESS = 300  # 限制处理数量，防止API限流


class StockScreener:
    def __init__(self):
        self.config = ScreenerConfig()
        self.fetch_time = datetime.now()
        self.stats = {
            '初始股票池': 0,
            '基础条件过滤后': 0,
            '行业筛选后': 0,
            '技术面筛选后': 0,
            '最终入选': 0
        }
        print(f"\n{'='*80}")
        print(f"A股股票智能筛选系统 v2.0".center(80))
        print(f"{'='*80}")
        print(f"数据获取时间: {self.fetch_time.strftime('%Y-%m-%d %H:%M:%S')}\n")
    
    def run(self) -> pd.DataFrame:
        """执行完整筛选流程"""
        # 1. 获取所有股票实时行情
        print("【阶段1】获取A股实时行情数据...")
        try:
            all_stocks = ak.stock_zh_a_spot_em()
            self.stats['初始股票池'] = len(all_stocks)
            print(f"  ✓ 成功获取 {len(all_stocks)} 支股票\n")
        except Exception as e:
            print(f"  ✗ 获取失败: {e}")
            return pd.DataFrame()
        
        # 2. 基础条件过滤
        print("【阶段2】基础条件筛选...")
        df = self._filter_basic(all_stocks)
        self.stats['基础条件过滤后'] = len(df)
        print(f"  ✓ 通过基础条件: {len(df)} 支\n")
        
        if df.empty:
            print("✗ 基础筛选后无剩余股票")
            return pd.DataFrame()
        
        # 3. 行业筛选
        print("【阶段3】行业优选...")
        df = self._filter_industry(df)
        self.stats['行业筛选后'] = len(df)
        print(f"  ✓ 优先行业股票: {len(df)} 支\n")
        
        # 4. 获取详细数据
        print("【阶段4】获取详细行情数据...")
        df_enriched = self._enrich_data(df)
        print(f"  ✓ 完成数据增强\n")
        
        # 5. 技术面筛选
        print("【阶段5】技术面筛选...")
        df = self._filter_technical(df_enriched)
        self.stats['技术面筛选后'] = len(df)
        print(f"  ✓ 通过技术面筛选: {len(df)} 支\n")
        
        # 6. 综合评分排序
        print("【阶段6】综合评分与排序...")
        df = self._score_and_rank(df)
        df = df.head(self.config.OUTPUT_COUNT)
        self.stats['最终入选'] = len(df)
        print(f"  ✓ 最终入选: {len(df)} 支\n")
        
        return df
    
    def _filter_basic(self, df: pd.DataFrame) -> pd.DataFrame:
        """基础条件过滤"""
        # 排除ST、退市股
        df = df[~df['名称'].str.contains('ST|退|N', na=False)]
        
        # 排除北交所（代码以4或8开头）
        df = df[~df['代码'].str.startswith(('4', '8'))]
        
        # 股价区间
        df = df[(df['最新价'] >= self.config.MIN_PRICE) & 
                (df['最新价'] <= self.config.MAX_PRICE)]
        
        # 成交额
        if '成交额' in df.columns:
            df = df[df['成交额'] >= self.config.MIN_VOLUME]
        
        # 总市值
        if '总市值' in df.columns:
            df = df[df['总市值'] >= self.config.MIN_MARKET_CAP]
        
        # PE、PB筛选
        if '市盈率-动态' in df.columns:
            df = df[(df['市盈率-动态'] >= self.config.MIN_PE) & 
                    (df['市盈率-动态'] <= self.config.MAX_PE)]
        
        if '市净率' in df.columns:
            df = df[(df['市净率'] >= self.config.MIN_PB) & 
                    (df['市净率'] <= self.config.MAX_PB)]
        
        # 换手率过滤（活跃度）
        if '换手率' in df.columns:
            df = df[df['换手率'] >= self.config.MIN_TURNOVER]
        
        return df.reset_index(drop=True)
    
    def _filter_industry(self, df: pd.DataFrame) -> pd.DataFrame:
        """行业筛选"""
        # 获取行业信息
        industries = []
        print(f"  正在获取行业信息...")
        
        for idx, row in df.head(self.config.MAX_STOCKS_TO_PROCESS).iterrows():
            try:
                symbol = row['代码']
                info = ak.stock_individual_info_em(symbol=symbol)
                industry_dict = dict(zip(info['item'], info['value']))
                industry = industry_dict.get('行业', '未知')
                list_date = industry_dict.get('上市时间', '')
                
                # 检查上市时间
                if list_date:
                    try:
                        list_dt = datetime.strptime(str(list_date), '%Y%m%d')
                        years = (datetime.now() - list_dt).days / 365
                        if years < self.config.MIN_LIST_YEARS:
                            continue
                    except:
                        pass
                
                # 检查行业
                if any(keyword in industry for keyword in self.config.PRIORITY_INDUSTRIES):
                    industries.append({
                        '代码': symbol,
                        '名称': row['名称'],
                        '行业': industry,
                        '最新价': row['最新价'],
                        '涨跌幅': row.get('涨跌幅', 0),
                        '换手率': row.get('换手率', 0),
                        '成交额': row.get('成交额', 0),
                        '总市值': row.get('总市值', 0),
                        '流通市值': row.get('流通市值', 0),
                        '市盈率': row.get('市盈率-动态', 0),
                        '市净率': row.get('市净率', 0),
                        '上市时间': list_date
                    })
                
                if idx % 50 == 0:
                    print(f"    进度: {idx}/{len(df)}, 已筛选: {len(industries)}")
                    
            except Exception as e:
                continue
        
        return pd.DataFrame(industries) if industries else pd.DataFrame()
    
    def _enrich_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """增强数据：获取历史行情"""
        results = []
        total = min(len(df), 100)  # 限制数量
        
        for idx, row in df.head(total).iterrows():
            try:
                symbol = row['代码']
                
                # 获取历史行情
                end_date = datetime.now().strftime('%Y%m%d')
                start_date = (datetime.now() - timedelta(days=300)).strftime('%Y%m%d')
                hist = ak.stock_zh_a_hist(symbol=symbol, period="daily", 
                                         start_date=start_date, end_date=end_date, adjust="qfq")
                
                if not hist.empty and len(hist) >= 60:
                    # 计算技术指标
                    current_price = row['最新价']
                    ma60 = hist['收盘'].tail(60).mean()
                    max_52w = hist['收盘'].tail(250).max() if len(hist) >= 250 else hist['收盘'].max()
                    
                    # 3个月涨跌幅
                    if len(hist) >= 60:
                        price_60d = hist.iloc[-60]['收盘']
                        return_3m = ((current_price - price_60d) / price_60d) * 100
                    else:
                        return_3m = 0
                    
                    # 52周回撤
                    drawdown_52w = ((max_52w - current_price) / max_52w) * 100
                    
                    row['60日均价'] = ma60
                    row['52周最高'] = max_52w
                    row['3个月涨跌幅'] = return_3m
                    row['52周回撤'] = drawdown_52w
                    
                    results.append(row.to_dict())
                
                if (idx + 1) % 20 == 0:
                    print(f"    进度: {idx+1}/{total}")
                    
            except Exception as e:
                continue
        
        return pd.DataFrame(results) if results else pd.DataFrame()
    
    def _filter_technical(self, df: pd.DataFrame) -> pd.DataFrame:
        """技术面筛选"""
        if df.empty:
            return df
        
        # 3个月涨跌幅
        if '3个月涨跌幅' in df.columns:
            df = df[(df['3个月涨跌幅'] >= self.config.MIN_3M_RETURN) & 
                    (df['3个月涨跌幅'] <= self.config.MAX_3M_RETURN)]
        
        # 52周回撤
        if '52周回撤' in df.columns:
            df = df[df['52周回撤'] <= self.config.MAX_52W_DRAWDOWN]
        
        # MA60支撑
        if '60日均价' in df.columns:
            df = df[df['最新价'] >= df['60日均价']]
        
        return df.reset_index(drop=True)
    
    def _score_and_rank(self, df: pd.DataFrame) -> pd.DataFrame:
        """综合评分"""
        if df.empty:
            return df
        
        # 评分模型：估值30%，技术面40%，流动性30%
        df['估值得分'] = 0
        df['技术得分'] = 0
        df['流动性得分'] = 0
        
        # 估值得分（PE越低越好）
        if '市盈率' in df.columns:
            pe_norm = (self.config.MAX_PE - df['市盈率']) / (self.config.MAX_PE - self.config.MIN_PE)
            df['估值得分'] = pe_norm.clip(0, 1) * 30
        
        # 技术得分（涨幅适中，回撤小）
        if '3个月涨跌幅' in df.columns:
            # 涨幅5-15%最佳
            return_score = 1 - abs(df['3个月涨跌幅'] - 5) / 15
            df['技术得分'] += return_score.clip(0, 1) * 20
        
        if '52周回撤' in df.columns:
            drawdown_score = 1 - df['52周回撤'] / self.config.MAX_52W_DRAWDOWN
            df['技术得分'] += drawdown_score.clip(0, 1) * 20
        
        # 流动性得分（成交额和换手率）
        if '成交额' in df.columns:
            volume_score = np.log10(df['成交额'] / 100_000_000).clip(0, 2) / 2
            df['流动性得分'] += volume_score * 15
        
        if '换手率' in df.columns:
            turnover_score = df['换手率'].clip(0, 10) / 10
            df['流动性得分'] += turnover_score * 15
        
        # 综合得分
        df['综合得分'] = df['估值得分'] + df['技术得分'] + df['流动性得分']
        
        # 生成推荐理由
        df['推荐理由'] = df.apply(self._generate_reason, axis=1)
        
        # 排序
        df = df.sort_values('综合得分', ascending=False)
        
        return df
    
    def _generate_reason(self, row) -> str:
        """生成推荐理由"""
        reasons = []
        
        # 行业
        if pd.notna(row.get('行业')):
            reasons.append(f"{row['行业']}")
        
        # 估值
        if row.get('市盈率', 0) < 15:
            reasons.append('低估值')
        elif row.get('市盈率', 0) < 20:
            reasons.append('估值合理')
        
        # 技术面
        return_3m = row.get('3个月涨跌幅', 0)
        if 0 <= return_3m <= 10:
            reasons.append('温和上涨')
        elif return_3m > 10:
            reasons.append('强势')
        
        # 回撤
        if row.get('52周回撤', 100) < 15:
            reasons.append('近高位')
        
        # 流动性
        if row.get('换手率', 0) > 3:
            reasons.append('活跃')
        
        return '，'.join(reasons) if reasons else '综合表现优秀'
    
    def print_summary(self):
        """打印筛选统计"""
        print(f"\n{'='*80}")
        print("筛选统计摘要".center(80))
        print(f"{'='*80}")
        for key, value in self.stats.items():
            print(f"  {key}: {value:>6} 支")
        
        if self.stats['初始股票池'] > 0:
            rate = (self.stats['最终入选'] / self.stats['初始股票池']) * 100
            print(f"  筛选率: {rate:>6.2f}%")
        print(f"{'='*80}\n")


class ReportGenerator:
    """报告生成类"""
    
    def __init__(self, df: pd.DataFrame, stats: dict, fetch_time: datetime):
        self.df = df
        self.stats = stats
        self.fetch_time = fetch_time
    
    def print_console(self):
        """控制台输出"""
        from tabulate import tabulate
        
        if self.df.empty:
            print("✗ 无筛选结果")
            return
        
        print(f"\n{'='*100}")
        print("A股智能筛选结果 - Top 20".center(100))
        print(f"{'='*100}\n")
        
        # 准备显示数据
        display_df = self.df.copy()
        
        # 选择核心字段
        cols = ['代码', '名称', '行业', '最新价', '市盈率', '市净率', 
                '3个月涨跌幅', '52周回撤', '换手率', '综合得分', '推荐理由']
        
        display_cols = [c for c in cols if c in display_df.columns]
        display_df = display_df[display_cols]
        
        # 格式化数值
        for col in display_df.columns:
            if col in ['最新价', '市盈率', '市净率']:
                display_df[col] = display_df[col].apply(lambda x: f"{x:.2f}" if pd.notna(x) else '-')
            elif col in ['3个月涨跌幅', '52周回撤', '换手率']:
                display_df[col] = display_df[col].apply(lambda x: f"{x:.1f}%" if pd.notna(x) else '-')
            elif col == '综合得分':
                display_df[col] = display_df[col].apply(lambda x: f"{x:.1f}" if pd.notna(x) else '-')
        
        print(tabulate(display_df, headers='keys', tablefmt='grid', showindex=False))
        
        # 风险提示
        print(f"\n{'='*100}")
        print("⚠️  风险提示".center(100))
        print(f"{'='*100}")
        print("  1. 本筛选结果仅供参考，不构成投资建议")
        print("  2. 股市有风险，投资需谨慎")
        print("  3. 历史业绩不代表未来表现")
        print("  4. 请结合个人研究与风险承受能力决策")
        print(f"{'='*100}\n")
    
    def export_excel(self, filename: str = None):
        """导出Excel"""
        if self.df.empty:
            print("✗ 无数据可导出")
            return
        
        if filename is None:
            filename = f"A股筛选结果_{self.fetch_time.strftime('%Y%m%d_%H%M%S')}.xlsx"
        
        try:
            with pd.ExcelWriter(filename, engine='openpyxl') as writer:
                # Sheet1: 筛选结果
                output_cols = ['代码', '名称', '行业', '最新价', '总市值', '市盈率', '市净率',
                              '3个月涨跌幅', '52周回撤', '换手率', '成交额', 
                              '综合得分', '推荐理由']
                
                out_df = self.df[[c for c in output_cols if c in self.df.columns]].copy()
                out_df.to_excel(writer, sheet_name='筛选结果', index=False)
                
                # Sheet2: 完整数据
                self.df.to_excel(writer, sheet_name='详细数据', index=False)
                
                # Sheet3: 参数说明
                params = pd.DataFrame({
                    '参数': ['股价区间', 'PE区间', 'PB区间', '3月涨跌幅', '52周回撤', '最小成交额', '最小市值'],
                    '标准': [
                        f"{ScreenerConfig.MIN_PRICE}-{ScreenerConfig.MAX_PRICE}元",
                        f"{ScreenerConfig.MIN_PE}-{ScreenerConfig.MAX_PE}",
                        f"{ScreenerConfig.MIN_PB}-{ScreenerConfig.MAX_PB}",
                        f"{ScreenerConfig.MIN_3M_RETURN}%~{ScreenerConfig.MAX_3M_RETURN}%",
                        f"<{ScreenerConfig.MAX_52W_DRAWDOWN}%",
                        f">{ScreenerConfig.MIN_VOLUME/100000000:.0f}亿元",
                        f">{ScreenerConfig.MIN_MARKET_CAP/100000000:.0f}亿元"
                    ]
                })
                params.to_excel(writer, sheet_name='筛选参数', index=False)
                
                # Sheet4: 统计
                stats_df = pd.DataFrame(list(self.stats.items()), columns=['统计项', '数量'])
                stats_df.to_excel(writer, sheet_name='筛选统计', index=False)
            
            print(f"✓ Excel已导出: {filename}")
        except Exception as e:
            print(f"✗ 导出失败: {e}")


def main():
    """主函数"""
    # 执行筛选
    screener = StockScreener()
    result_df = screener.run()
    screener.print_summary()
    
    if result_df.empty:
        print("✗ 未找到符合条件的股票，建议调整筛选参数")
        return
    
    # 生成报告
    reporter = ReportGenerator(result_df, screener.stats, screener.fetch_time)
    reporter.print_console()
    reporter.export_excel()
    
    print("✓ 筛选任务完成！\n")


if __name__ == "__main__":
    main()
